from fastapi import FastAPI
from pydantic import BaseModel
from openai import OpenAI
import os

import server_config

app = FastAPI(title="NLU Processing Service")

# 配置参数（实际使用时应从环境变量或配置文件中读取）
OPENAI_API_BASE = os.getenv("OPENAI_API_BASE", "http://0.0.0.0:10085/v1")
MODEL_DOMAIN = os.getenv("MODEL_DOMAIN", "lora_domain")
MODEL_SEMANTIC = os.getenv("MODEL_SEMANTIC", "lora_semantic")
MODEL_REWRITE = os.getenv("MODEL_REWRITE", "base_model")

# 系统提示词配置
system_content_rewrite = server_config.SYSTEM_CONTENT_REWRITE
system_message = [{"role": "system", "content": system_content_rewrite}]

system_content_segment = server_config.SYSTEM_CONTENT_SEGMENT

# 初始化OpenAI客户端
client = OpenAI(
    api_key="EMPTY",  # vLLM不需要实际API密钥
    base_url=OPENAI_API_BASE
)

class rewriteRequest(BaseModel):
    query_history: list
    query: str

class rewriteResponse(BaseModel):
    result: dict

class segRequest(BaseModel):
    query: str

class segResponse(BaseModel):
    result: dict

@app.post("/segment", response_model=segResponse)
async def segment(request: segRequest):
    query = request.query
    chat_response = client.chat.completions.create(
        model="base_model",
        messages=[
            {"role": "system", "content": system_content_segment},
            {"role": "user", "content": query}
        ],
        max_tokens=128,
        extra_body={
            "chat_template_kwargs": {"enable_thinking": False},
        },
    )
    res = chat_response.choices[0].message.content
    result = { "query_raw": query, "query_new": res}
    return segResponse(
        result=result
    )

@app.post("/rewrite", response_model=rewriteResponse)
async def rewrite(request: rewriteRequest):
    query = request.query
    messages_now = [{"role": "user", "content": query }]
    
    query_history = request.query_history
    print(query_history)
    messages_history = []
    if len(query_history) > 0:
        for message in query_history:
            if len(message) == 1:
                messages_history.append({"role": "user", "content": message[0]})
            if len(message) == 2:
                messages_history.append({"role": "user", "content": message[0]})
                messages_history.append({"role": "assistant", "content": message[1]})

    messages = system_message + messages_history + messages_now

    response = client.chat.completions.create(
        model=MODEL_REWRITE,
        messages=messages,
        temperature=0.2,
        max_tokens=128,
        extra_body={"chat_template_kwargs": {"enable_thinking": False}}
    )

    query_rewritten = response.choices[0].message.content

    result = { "query_raw": query, "query_new": query_rewritten}

    return rewriteResponse(
        result=result
    )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=10086)